#include "VideoEncoder.h"
#include "common.h"

CVideoEncoder::CVideoEncoder(EncodeDataNotify *pNotify)
	: m_pNotify(pNotify)
	, m_hEncodeHandle(NULL)
	, m_bExit(false)
	, m_pH264Buf(NULL)
	, m_nSrcWidth(0)
	, m_nSrcHeight(0)
	, m_pOutFile(NULL)
	, m_nBuffCount(0)
{
	//m_pOutFile = fopen("d:\\test.h264", "wb");

	StartThread();
}

CVideoEncoder::~CVideoEncoder(void)
{
	m_bExit = true;

	WaitForStop();

	CloseEncoder();

	while (m_lstYUVData.size())
	{
		YUVDataBuff *pYUV = NULL;

		pYUV = (YUVDataBuff *)m_lstYUVData.front();
		m_lstYUVData.pop_front();

		if (pYUV)
			ClearRGBBuf(pYUV);
	}

	while (m_lstFree.size())
	{
		YUVDataBuff *pYUV = NULL;

		pYUV = (YUVDataBuff *)m_lstFree.front();
		m_lstFree.pop_front();

		if (pYUV)
			ClearRGBBuf(pYUV);
	}

	if (m_pH264Buf)
	{
		delete [] m_pH264Buf;
		m_pH264Buf = NULL;
	}
}

void CVideoEncoder::PutYUVData(const uint8_t * const pFrame[], int nWidth, int nHeight, unsigned long ulTimestamp)
{
	YUVDataBuff *pYUVData = GetYuvBuff();
	if (pYUVData)
	{
		if (pYUVData->nDataLen == 0)
		{
			pYUVData->nDataLen = nWidth * nHeight * 3 / 2;
			pYUVData->pData = new char[pYUVData->nDataLen];				
		}

		pYUVData->nWidth = nWidth;
		pYUVData->nHeigth = nHeight;
		pYUVData->ulTimestamp = ulTimestamp;

		if (pYUVData->pData)
		{
			int nLen = 0;
			memcpy(pYUVData->pData, pFrame[0], nWidth * nHeight);

			nLen += nWidth * nHeight;
			memcpy(pYUVData->pData + nLen, pFrame[1], nWidth * nHeight / 4);

			nLen += nWidth * nHeight / 4;
			memcpy(pYUVData->pData + nLen, pFrame[2], nWidth * nHeight / 4);

			CAutoLock autoLock(&m_lockYUVData);
			m_lstYUVData.push_back(pYUVData);
		}
	}
}

void CVideoEncoder::ClearRGBBuf(YUVDataBuff *pYUV)
{
	if (pYUV)
	{
		if (pYUV->pData)
		{
			delete [] pYUV->pData;
			pYUV->pData = NULL;
		}

		delete pYUV;
		pYUV = NULL;
	}
}

bool CVideoEncoder::OpenEncoder(const EncoderParam &encoderParam)
{
	CloseEncoder();

	if (m_pH264Buf)
	{
		delete [] m_pH264Buf;
		m_pH264Buf = NULL;
	}

	m_nSrcWidth = encoderParam.nWidth;
	m_nSrcHeight = encoderParam.nHeight;	
	
	m_hEncodeHandle = H264enc_open(encoderParam.nWidth, encoderParam.nHeight, encoderParam.nFrame, encoderParam.nQuant, encoderParam.nBitRate, encoderParam.nQMin, encoderParam.nQMax, encoderParam.nThreadNum);
	if (m_hEncodeHandle != NULL)
	{
		m_param = encoderParam;

		m_nBufLen = encoderParam.nWidth * encoderParam.nHeight;
		m_pH264Buf = new char[m_nBufLen];
	}	

	return m_hEncodeHandle != NULL;
}

void CVideoEncoder::CloseEncoder(void)
{
	if (m_hEncodeHandle)
	{
		H264enc_close(m_hEncodeHandle);
		m_hEncodeHandle = NULL;
	}
}

int CVideoEncoder::Encode(YUVDataBuff *pYUV)
{
	if (pYUV->nWidth != m_nSrcWidth || pYUV->nHeigth != m_nSrcHeight)
	{
		CloseEncoder();

		m_nSrcWidth = pYUV->nWidth;
		m_nSrcHeight = pYUV->nHeigth;

		if (!OpenEncoder(m_param))
			return -1;
	}

	if (m_hEncodeHandle == NULL)
	{
		OpenEncoder(m_param);
		if (m_hEncodeHandle == NULL)
			return -1;
	}

	int nKeyFrame(0),nCurrentLen(0);	

	unsigned char *pYV12[3];
	pYV12[0] = (unsigned char *)pYUV->pData;

	nCurrentLen += m_nSrcWidth * m_nSrcHeight;
	pYV12[1] = (unsigned char *)(pYUV->pData + nCurrentLen);

	nCurrentLen += m_nSrcWidth * m_nSrcHeight / 4;
	pYV12[2] = (unsigned char *)(pYUV->pData + nCurrentLen);

	int iRet = H264enc_encode(m_hEncodeHandle, pYV12, (unsigned char*)m_pH264Buf, m_nBufLen, &nKeyFrame);
	if (iRet > 0 && m_pNotify)
	{
		H264VideoParam h264VideoFrame;
		h264VideoFrame.pBuf = m_pH264Buf;
		h264VideoFrame.nDataLen = iRet;
		h264VideoFrame.nWidth = m_param.nWidth;
		h264VideoFrame.nHeight = m_param.nHeight;
		h264VideoFrame.nFrameRate = m_param.nFrame;
		h264VideoFrame.nIsKeyFrame = nKeyFrame;
		h264VideoFrame.ulTimestamp = pYUV->ulTimestamp;

		//fwrite(m_pH264Buf, 1, iRet, m_pOutFile);

		m_pNotify->OnEncodeData(&h264VideoFrame);

		memset(m_pH264Buf, 0x00, m_nBufLen);
	}

	return iRet;
}

YUVDataBuff * CVideoEncoder::GetYuvBuff(void)
{
	YUVDataBuff *pYuv = NULL;

	if ((int)m_lstFree.size() > 0)
	{
		pYuv = (YUVDataBuff *)m_lstFree.front();
		m_lstFree.pop_front();

		return pYuv;
	}
	else
	{
		if (m_nBuffCount < 50)
		{
			pYuv = new YUVDataBuff;
			if (pYuv)
			{
				pYuv->nDataLen = 0;
				pYuv->pData = NULL;

				m_nBuffCount++;
			}
		}		
	}

	return pYuv;
}

void CVideoEncoder::ThreadProcMain(void)
{
	while (!m_bExit)
	{
		while (m_lstYUVData.size() && !m_bExit)
		{
			YUVDataBuff *pYUV = NULL;

			{
				CAutoLock autoLock(&m_lockYUVData);
				pYUV = (YUVDataBuff *)m_lstYUVData.front();
				m_lstYUVData.pop_front();
			}

			if (pYUV)
				Encode(pYUV);

			m_lstFree.push_back(pYUV);

			//ClearRGBBuf(pYUV);
		}

		my_sleep(5);
	}
}
